# -*- coding: UTF-8 -*-
# @IDE     : VScode
# @File   : zhipu_llm.py
# @Time   : 2024/05/26 09:37:41
# @Author : zhonggc

from typing import List, Any, Dict, Iterator, Mapping, Optional
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models import LLM
from langchain_core.outputs import GenerationChunk
from zhipuai import ZhipuAI


class ZhiPu(LLM):
    """质谱AI模型

    Args:
        LLM (_type_): _description_
    """

    # 默认选用 glm-4 模型
    model: str = "glm-4"
    # 温度系数
    temperature: float = 0.1
    # API_Key
    api_key: str = None

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        """模型生成回答的逻辑方法

        Args:
            prompt (str): _description_
            stop (Optional[List[str]], optional): _description_. Defaults to None.
            run_manager (Optional[CallbackManagerForLLMRun], optional): _description_. Defaults to None.

        Returns:
            str: _description_
        """

        if self.api_key is None:
            raise ValueError("API_Key is required.")
        client = ZhipuAI(api_key=self.api_key)  # 请填写您自己的APIKey
        response = client.chat.completions.create(
            model="glm-4",  # 填写需要调用的模型名称
            messages=[
                {"role": "user", "content": prompt},
            ],
            stream=False,  # 是否流式处理
        )
        if len(response.choices) > 0:
            return response.choices[0].message.content
        return "generate answer error."
    

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        """标识参数

        Returns:
            Dict[str, Any]: _description_
        """
        return {"model_name": "ZhiPu-GLM"}

    @property
    def _llm_type(self) -> str:
        """获取此聊天模型使用的语言模型类型。仅用于日志记录目的。

        Returns:
            str: _description_
        """
        return "ZhiPu-GLM"
